package com.addzero.web.ui.hooks.tree

import androidx.compose.foundation.layout.*
import androidx.compose.material.icons.Icons
import androidx.compose.material.icons.automirrored.filled.ArrowRight
import androidx.compose.material.icons.filled.ArrowDropDown
import androidx.compose.material3.Icon
import androidx.compose.material3.IconButton
import androidx.compose.material3.TriStateCheckbox
import androidx.compose.runtime.Composable
import androidx.compose.runtime.getValue
import androidx.compose.runtime.mutableStateOf
import androidx.compose.runtime.setValue
import androidx.compose.ui.Alignment
import androidx.compose.ui.Modifier
import androidx.compose.ui.state.ToggleableState
import androidx.compose.ui.unit.dp
import com.addzero.common.kt_util.isNotEmpty
import com.addzero.web.ui.hooks.UseHook

/**
 * 节点状态管理器 - 处理节点选中和半选中状态的计算
 */
class NodeState<T>(
    private val getIdFun: (T) -> Any,
    private val getChildren: (T) -> List<T>
) {
    // 完全选中的节点ID集合
    var selectedNodes by mutableStateOf(setOf<Any>())
        private set
    
    // 半选中（部分子节点选中）的节点ID集合
    var indeterminateNodes by mutableStateOf(setOf<Any>())
        private set
    
    // 切换节点的选中状态
    fun toggleNodeSelection(node: T, allNodes: List<T>) {
        val nodeId = getIdFun(node)
        
        // 如果节点已选中，则取消选中它和它的所有子节点
        if (isNodeSelected(nodeId)) {
            // 找出所有子节点ID
            val childrenIds = collectAllChildrenIds(node)
            // 移除该节点及其所有子节点
            selectedNodes = selectedNodes - setOf(nodeId) - childrenIds
        } else {
            // 如果节点未选中，则选中它和它的所有子节点
            val childrenIds = collectAllChildrenIds(node)
            // 添加该节点及其所有子节点
            selectedNodes = selectedNodes + setOf(nodeId) + childrenIds
        }
        
        // 重新计算所有半选中状态
        recalculateIndeterminateStates(allNodes)
    }
    
    // 检查节点是否被选中
    fun isNodeSelected(nodeId: Any): Boolean {
        return selectedNodes.contains(nodeId)
    }
    
    // 检查节点是否为半选中状态
    fun isNodeIndeterminate(nodeId: Any): Boolean {
        return indeterminateNodes.contains(nodeId)
    }
    
    // 递归收集所有子节点的ID
    private fun collectAllChildrenIds(node: T): Set<Any> {
        val result = mutableSetOf<Any>()
        val children = getChildren(node)
        
        for (child in children) {
            val childId = getIdFun(child)
            result.add(childId)
            // 递归收集子节点的子节点
            result.addAll(collectAllChildrenIds(child))
        }
        
        return result
    }
    
    // 重新计算所有半选中状态
    private fun recalculateIndeterminateStates(allNodes: List<T>) {
        val newIndeterminateNodes = mutableSetOf<Any>()
        
        // 构建节点ID到节点的映射，以便快速查找父节点
        val nodesMap = mutableMapOf<Any, T>()
        val childToParentMap = mutableMapOf<Any, Any>()
        
        // 填充映射
        fun mapNodes(nodes: List<T>, parentId: Any? = null) {
            for (node in nodes) {
                val nodeId = getIdFun(node)
                nodesMap[nodeId] = node
                
                if (parentId != null) {
                    childToParentMap[nodeId] = parentId
                }
                
                mapNodes(getChildren(node), nodeId)
            }
        }
        
        mapNodes(allNodes)
        
        // 自底向上计算半选中状态
        // 遍历所有节点
        for (node in allNodes) {
            calculateIndeterminateState(node, newIndeterminateNodes, nodesMap, childToParentMap)
        }
        
        indeterminateNodes = newIndeterminateNodes
    }
    
    // 计算单个节点的半选中状态
    private fun calculateIndeterminateState(
        node: T, 
        indeterminateNodes: MutableSet<Any>,
        nodesMap: Map<Any, T>,
        childToParentMap: Map<Any, Any>
    ) {
        val nodeId = getIdFun(node)
        val children = getChildren(node)
        
        if (children.isEmpty()) {
            // 叶子节点不会是半选中状态
            return
        }
        
        // 计算子节点的选中情况
        val childrenIds = children.map { getIdFun(it) }
        val allChildrenSelected = childrenIds.all { selectedNodes.contains(it) }
        val noChildrenSelected = childrenIds.none { selectedNodes.contains(it) || indeterminateNodes.contains(it) }
        
        // 如果部分子节点被选中（不是全部且不是零），则该节点为半选中状态
        if (!allChildrenSelected && !noChildrenSelected) {
            indeterminateNodes.add(nodeId)
            
            // 向上传播半选中状态
            var currentNodeId = nodeId
            while (childToParentMap.containsKey(currentNodeId)) {
                val parentId = childToParentMap[currentNodeId]!!
                indeterminateNodes.add(parentId)
                currentNodeId = parentId
            }
        }
    }
}

class UseTreeNodeItem<T>(
    val editmodeFlag: Boolean = false,
    val getLabel: (T) -> String,
    val getChildren: (T) -> List<T>,
    val getNodeTypeFun: (T) -> NodeTypes,
    val getIdFun: (T) -> Any,
    var customRender: (@Composable (T) -> Unit) = {
        defaultTreeNodeRender(it!!, getNodeTypeFun, getLabel = getLabel)
    },
) : UseHook<UseTreeNodeItem<T>> {
    var node by mutableStateOf(null as T?)
    var level by mutableStateOf(0)

    // 节点是否展开
    var expanded by mutableStateOf(false)

    // 安全地获取子节点，如果node为null则返回空列表
    private fun safeGetChildren(): List<T> {
        return if (node != null) {
            try {
                getChildren(node!!)
            } catch (e: Exception) {
                emptyList()
            }
        } else {
            emptyList()
        }
    }

    // 计算子节点列表
    var children by mutableStateOf(emptyList<T>())

    // 是否是叶子节点
    var leafFlag by mutableStateOf(true)

    // 是否选中
    var selected by mutableStateOf(false)
    
    // 节点状态管理器 - 懒初始化，仅在编辑模式下使用
    private val nodeState by lazy { 
        NodeState<T>(getIdFun, getChildren)
    }

    // 更新节点数据时同时更新子节点列表
    internal fun updateNode(newNode: T?) {
        node = newNode
        if (newNode != null) {
            children = safeGetChildren()
            leafFlag = children.isEmpty()
        } else {
            children = emptyList()
            leafFlag = true
        }
    }

    override val render: @Composable () -> Unit
        get() = {
            if (node == null) {
                // 如果节点为空，不渲染任何内容
            } else {
                val nodeId = getIdFun(node!!)
                val that = this
                
                // 如果在编辑模式下，获取当前节点的状态
                val isSelected = if (editmodeFlag) nodeState.isNodeSelected(nodeId) else selected
                val isIndeterminate = if (editmodeFlag) nodeState.isNodeIndeterminate(nodeId) else false

                Column(modifier = modifier) {
                    Row(
                        modifier = Modifier.fillMaxWidth().padding(start = (level * 16).dp)
                            .padding(vertical = 4.dp),
                        verticalAlignment = Alignment.CenterVertically
                    ) {
                        // 显示展开/折叠箭头（如果有子节点）
                        if (!leafFlag) {
                            IconButton(
                                onClick = { expanded = !expanded },
                                modifier = Modifier.size(24.dp)
                            ) {
                                Icon(
                                    imageVector = if (expanded) Icons.Default.ArrowDropDown else Icons.AutoMirrored.Filled.ArrowRight,
                                    contentDescription = if (expanded) "折叠" else "展开"
                                )
                            }
                        } else {
                            // 为叶子节点添加空白，保持对齐
                            Spacer(modifier = Modifier.width(24.dp))
                        }
                        
                        // 编辑模式的三态切换
                        if (editmodeFlag) {
                            // 获取三态复选框的状态
                            val checkboxState = when {
                                isSelected -> ToggleableState.On
                                isIndeterminate -> ToggleableState.Indeterminate
                                else -> ToggleableState.Off
                            }
                            
                            TriStateCheckbox(
                                state = checkboxState,
                                onClick = {
                                    // 切换节点选中状态
                                    // 注意：理想情况下，我们应该从根节点获取所有节点
                                    // 这里简化处理，仅使用当前节点和其子节点
                                    node?.let { 
                                        val allNodes = listOf(it) + children
                                        nodeState.toggleNodeSelection(it, allNodes)
                                    }
                                }
                            )
                        }

                        // 根据节点类型和自定义渲染函数来渲染内容
                        customRender(node!!)
                    }

                    // 如果当前节点展开且有子节点，递归渲染子节点
                    if (expanded && children.isNotEmpty()) {
                        children.forEach { childNode ->
                            UseTreeNodeItem(
                                editmodeFlag,
                                getLabel,
                                getChildren,
                                getNodeTypeFun,
                                getIdFun,
                                customRender
                            ).apply {
                                updateNode(childNode)
                                level = that.level + 1 // 自动计算：子节点的level是父节点level + 1
                                expanded = false // 子节点默认折叠，让用户手动展开
                            }.render {}
                        }
                    }
                }
            }
        }
}
